"""Support for LSA on a corpus of text via random accumulation vectors.
"""
import os, sys, string
from nltk.corpus import wordnet as wn
from rav import rav
import stops

# -----------------------------------------------------------------------------
class LSA:
   def __init__(self, seed=-1, max_vectors=50000, dim=2048):
      self.ravhash = rav.rav_init(seed, max_vectors, dim)
      if not self.ravhash:
         raise Exception('failed to init ravhash')
      self.stops = stops.stopwords
      self.half_window_size = 2
      self.min_word_length = 3
      self.remove_chars = string.digits
      self.translate_chars = string.maketrans(\
        '!"#$%&()*+,-./:;<=>@[\]^_`{|}~\'\n\t',
        '                                 ')
      #self.translate_chars = string.maketrans(\
        #'!"#$%&()*+,-./:;<=>?@[\]^_`{|}~\'\n\t',
        #'                                  ')
   def __del__(self):
      try:
         rav.rav_destroy(self.ravhash)
      except:
         pass
   def get_words(self, text):
      if text is None:
         return []
      text = text.replace("?"," ")
      # make sure we've got plain ascii
      text = text.encode('ascii','replace')
      # turn puntuation into spaces
      text = text.translate(self.translate_chars, self.remove_chars)
      # split based on whitespace
      words = text.lower().split()
      # morphological stemming
      for i,word in enumerate(words):
         root = wn.morphy(word)
         if root and root!=word:
            words[i] = root
      # remove words that are too short or on stop-list
      words = [x for x in words if (len(x)>=self.min_word_length and not \
        x in self.stops)]
      return words
   def get_ids(self, words, auto_add=True):
      ids = []
      for word in words:
         id = self.lookup(word)
         if id == -1:
            if auto_add: # not already in hash
               id = self.new(word)
            else:
               raise KeyError(word)
         ids.append(id)
      return ids
   def add_text(self, text):
      # slide window across text, accumulating each word in range -k <-> +k
      # (where k is self.half_window_size) and adding normalized sum to window
      # center's context.
      hsz = self.half_window_size
      center = hsz
      ids = self.get_ids(self.get_words(text))
      window = [None]*(hsz*2+1)
      for id in ids+[None]*hsz:
         window = window[1:]+[id]
         if window[center] is not None:
            self.accumulate_env(window[0:center]+window[center+1:], \
             window[center])
   def add_file(self, fname, show_progress=False, line_limit=0):
      with file(fname) as f:
         i=0
         for line in f:
            i+=1
            self.add_text(line.strip())
            if show_progress and (i%100)==0:
               sys.stderr.write('.')
            if line_limit and i>=line_limit:
               break
   def find_closest(self, text, n=10, tag=None, already_a_word=False):
      if already_a_word:
         ids = self.get_ids([text], auto_add=False)
      else:
         ids = self.get_ids(self.get_words(text), auto_add=False)
      tmp = self.tmp()
      if len(ids)==1:
         if tag is None:
            tmp = ids[0]
         else:
            self.convolve_env(ids, tag, tmp)
      else:
         if tag is None:
            self.accumulate_env(ids)
         else:
            self.convolve_env(ids, tag, tmp)
      rav.rav_closest(self.ravhash, tmp, n, 1)
      return self.get_closest_results(n)
   def get_closest_results(self, n):
      res = []
      for type in (rav.RAV_WORD, rav.RAV_FEATURE, rav.RAV_VALUE):
         for i in range(n):
            res.append((type,
             rav.rav_closest_idx_n(self.ravhash, type, i),\
             rav.rav_closest_name_n(self.ravhash, type, i),\
             rav.rav_closest_score_n(self.ravhash, type, i)))
      return res
   def print_closest(self, target, n=10, precision=2):
      fmt = '%%4.%df' % precision
      try:
         for type, idx, name, score in self.find_closest(target, n):
            if idx != -1:
               print '  ', name, fmt % score
      except KeyError, e:
         print "%s not found" % e
   def print_cosine(self, t1, t2):
      idx1,idx2 = self.lookup(t1), self.lookup(t2)
      if idx1<0 or idx2<0:
         print "'%s' or '%s' not found" % (t1, t2)
      print '%s %s : %4.3f' % (t1, t2, rav.rav_cosine(self.ravhash, idx1, idx2))
   def lookup(self, word):
      return rav.rav_lookup(self.ravhash, word, -1)
   def new(self, word):
      id = rav.rav_new(self.ravhash, word, -1)
      if id==-1:
         raise Exception("out of space")
      return id
   def clr_ctx(self, idx):
      rav.rav_clr_ctx(self.ravhash, idx)
   def accumulate_env(self, others, target=None):
      self.clr_ctx(self.tmp())
      for id in others:
         if id is not None:
            rav.rav_acc(self.ravhash, id)
      if target is not None:
         rav.rav_add_acc(self.ravhash, target)
      elif len(others)>1:
         rav.rav_normalize(self.ravhash, self.tmp())
   def convolve_env(self, others, tag, tmp, target=None):
      self.clr_ctx(tmp)
      for id in others:
         rav.rav_cconv_add(self.ravhash, tmp, tag, id)
      if target is not None:
         rav.rav_add_acc(self.ravhash, target)
      else:
         rav.rav_normalize(self.ravhash, tmp)
   def tmp(self):
      return rav.rav_get_tmp_idx(self.ravhash)
   def normalize(self, idx=-1):
      rav.rav_normalize(self.ravhash, idx)
   def delete(self, nm):
      return rav.rav_delete(self.ravhash, nm)
   def save(self, filename):
      with file(filename, 'wb+') as f:
         # pass file descriptor to c to write rav data
         if rav.rav_save(self.ravhash, f.fileno()):
            print('error saving rav hash')
         f.flush()
         os.fsync(f.fileno())
   def load(self, filename):
      with file(filename, 'rb') as f:
         # pass file descriptor to c to read rav data
         if rav.rav_load(self.ravhash, f.fileno()):
            sys.exit('error loading rav hash')
         print '(%d)' % self.tmp()
         self.onload()
   def dumppca(self, n=0):
      sys.stderr.write("dumping matrix for pca..")
      rav.rav_dump(self.ravhash, "dump.dat", n, 0)
      rav.rav_dump_names(self.ravhash, "dump.py", n, 0)
      sys.stderr.write("done.\n")
   def onload(self):
      pass
   def interact(self):
      """<word>          : find default number of closest neighbors to <word>
         <word> <N>      : find <N> closest neighbors to <word>
         <word1> <word2> : find cosine between words"""
      try:
         while True:
            print '>',
            t = raw_input()
            tok = t.split()
            if len(tok)==2:
               try:
                  n = int(tok[1].strip())
                  self.print_closest(tok[0].strip(), n)
               except ValueError:
                  self.rav.print_cosine(tok[0].strip(), tok[1].strip())
            elif len(tok)==1:
               if tok[0]=='dumppca':
                  self.dumppca()
               elif tok[0] in ('quit', 'exit'):
                  break
               else:
                  self.print_closest(tok[0].strip())
      except (KeyboardInterrupt, EOFError):
         pass

# -----------------------------------------------------------------------------
# example usage
# -----------------------------------------------------------------------------
if __name__ == '__main__':
   def test1():
      r = LSA()
      sys.stderr.write('processing text..')
      r.add_file('wiki.txt')
      sys.stderr.write('done\n')
      r.normalize()
      print '%d words' % rav.rav_num_unique(r.ravhash)
      print 'lookup "baseball":'
      r.print_closest('baseball', 10)
      r.print_cosine('baseball','game')
      r.dumppca() # dump info for principle component analysis
      sys.stderr.write('saving matrix..')
      r.save('wiki.rav') # save semantic space as matrix
      sys.stderr.write('done\n')
   def test2():
      r = LSA()
      r.load('wiki.rav')
      r.print_closest('baseball', 10)
      rav.rav_delete(r.ravhash, 'baseball')
      r.print_closest('baseball')
      r.print_closest('league')
   test1()
